import torch 
import torch.nn as nn 

import matplotlib.pyplot as plt 
import torchvision

image_size = 64
channels_img = 1
channels_noise = 256
features_g = 16

class Generator(nn.Module):
  def __init__(self, channels_noise, channels_img, features_g):
    super(Generator, self).__init__()

    self.net = nn.Sequential(
        # N x channels_noise x 1 x 1
        nn.ConvTranspose2d(channels_noise, features_g*16, kernel_size=4, stride=1, padding=0),
        nn.BatchNorm2d(features_g*16),
        nn.LeakyReLU(0.2),

        # N x features_g*16 x 4 x 4
        nn.ConvTranspose2d(features_g*16, features_g*8, kernel_size=4, stride=2, padding=1),
        nn.BatchNorm2d(features_g*8),
        nn.LeakyReLU(0.2),

        # N x features*8 x 8 x 8
        nn.ConvTranspose2d(features_g*8, features_g*4, kernel_size=4, stride=2, padding=1),
        nn.BatchNorm2d(features_g*4),
        nn.LeakyReLU(0.2),

        # N x features*4 x 16 x 16
        nn.ConvTranspose2d(features_g*4, features_g*2, kernel_size=4, stride=2, padding=1),
        nn.BatchNorm2d(features_g*2),
        nn.LeakyReLU(0.2),

        # N x features*2 x 32 x 32
        nn.ConvTranspose2d(features_g*2, channels_img, kernel_size=4, stride=2, padding=1),
        nn.Tanh()
    )

  def forward(self, x):
    return self.net(x)



if __name__ == "__main__":
    # print(torchvision.__version__)
    gan_model = Generator(channels_noise, channels_img, features_g)
    n_samples = 25
    latent_dim = 100 
    # latent_vec = torch.rand(latent_dim * n_samples)
    # latent_vec = latent_vec.view((n_samples, latent_dim))

    fixed_noise = torch.randn(n_samples, channels_noise, 1, 1)

    gan_model.load_state_dict(torch.load("./gan_model.pkl", map_location="cpu"))
    fake_image = gan_model(fixed_noise)
    # mean_v = torch.mean(fake_image)
    # fake_image[fake_image > mean_v] = 1
    # fake_image[fake_image < mean_v] = 0
    # print(fake_image[0])
    # fake_image = (fake_image + 1) / 2 
    # print(fake_image[0])
    print(fake_image.shape)

    for i in range(25):
        plt.subplot(5, 5, 1 + i)
        plt.axis("off")
        plt.imshow(fake_image[i].squeeze(0).detach().numpy(), cmap ='gray')
    plt.show()